# -*- encoding: utf-8 -*-
'''
@File     : server.py
@Project  : pipecoco
@Date     : 2021/8/13
@Author   : Zhang Jinyang
@Contact  : zhang-jy@sjtu.edu.cn
'''
 
from comm.pipecoco_server import *
from net.alexnet.loader import load_alexnet
from net.vgg16.loader import load_vgg_small, load_vgg_large
from net.resnet.loader import load_resnet34

# 创建PipeCoCo服务端示例
class server_test(pipecoco_server):

    def __init__(self, host):
        pipecoco_server.__init__(self, host)

    def load_net(self, model_name):

        if model_name == 'alexnet':
            network = load_alexnet('./models/alexnet.ckpt')

        elif model_name == 'vgg_small':
            network = load_vgg_small('./models/vgg_small.ckpt')

        elif model_name == 'vgg_large':
            network = load_vgg_large('./models/vgg_large.ckpt')

        elif model_name == ('resnet34'):
            network = load_resnet34('./models/resnet34.ckpt')

        return network


if __name__ == '__main__':

    # host改成真实IP
    host = "localhost:5000"
    model_path = "./models/"
    server = server_test(host)
    server.serve()